__author__="YR"
__date__ ="$Jul 9, 2013 11:56:13 PM$"

#show ground truth for BEHAVE data set

import os
import xml.etree.ElementTree as ET
import numpy as np
from skimage import io
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Define input image sequence
IMGSEQ_PATH = '../0-11200'
IMGSEQ_PREFIX = '000'
IMGSEQ_NDIGITS = 5
IMGSEQ_FTYPE = 'jpg'
START_F = 0 # default 0
END_F = 100 # last frame is 100
PLAYBACK_FPS = 1 # in ms, default 50

sequence0Tree = ET.parse('../1-11200.xgtf')
sequence0Root = sequence0Tree.getroot()

numOfFrames = 11201
##numOfFrames = len(os.walk(IMGSEQ_PATH).next()[2])
numOfPerson = 0
for person in sequence0Root[1][0].iter('{http://lamp.cfar.umd.edu/viper#}object'):
    numOfPerson += 1

frames = [[[]]]
for i in range(numOfFrames):
    singleframe = "frame"+str(i)
    singleframe = [[8, 0, 0, 0, 0]]
    frames.append(singleframe)

frames.pop(0)

#'{http://lamp.cfar.umd.edu/viper#}object'
person0 = sequence0Root[1][0][1]
person1 = sequence0Root[1][0][2]
person2 = sequence0Root[1][0][3]
person3 = sequence0Root[1][0][4]
person4 = sequence0Root[1][0][5]
person5 = sequence0Root[1][0][6]


for i in range(numOfPerson):
    curPerson = sequence0Root[1][0][i+1]

    for boxes in curPerson.iter('{http://lamp.cfar.umd.edu/viperdata#}bbox'):
        framespan = boxes.attrib["framespan"].split(':')
        box_cur = []
        box_cur.append(i)
        box_cur.append(boxes.attrib["height"])
        box_cur.append(boxes.attrib["width"])
        box_cur.append(boxes.attrib["x"])
        box_cur.append(boxes.attrib["y"])

        framespanstart = int(framespan[0])
        framespanend = int(framespan[1])
        for j in range(framespanstart-1, framespanend):
            frames[j].append(box_cur)

for i in range(numOfFrames):
    frames[i].pop(0)

# play frames [start_f : end_f] not including end_f
def imgseg_showbox (imgseq_path, imgseq_prefix, imgseq_ndigits, imgseq_ftype, start_f, end_f, playback_fps, frames):
  if start_f < 0:
    print 'imgseq_play Error: start_f = %d < 0' % (start_f)
    return False
  if end_f < start_f:
    print 'imgseq_play Error: end_f = %d < start_f %d' % (end_f, start_f)
    return False
  if not os.path.exists(imgseq_path):
    print 'imgseq_play Error: %s does not exist' % (imgseq_path)
  if imgseq_ndigits < 0 or imgseq_ndigits > 9: #30:
    print 'imgseq_play Error: imgseq_ndigits = %d' % (imgseq_ndigits)
    return False
  n_f = end_f - start_f

  # load first frame to obtain width, height
  digits_format = '%0'+str(imgseq_ndigits)+'d'
  fdigs = digits_format % (start_f)
  img_filename = '%s%s.%s' % (imgseq_prefix, fdigs, imgseq_ftype)
  img_fullpath = imgseq_path+'/'+img_filename
  try:
    image = io.imread (img_fullpath)
  except IOError:
    print 'imgseg_play Error in reading %s' % (img_fullpath)
    return False

  fig = plt.figure()
  ax = fig.add_subplot(111, aspect='equal', autoscale_on=True)

  # Note this initial setup won't show, it will be overwritten by updatefig.
  plt_im = plt.imshow (image, aspect='equal')
  # However the following init to np.zero array with double values yield error!
  #plt_im = plt.imshow (np.zeros(image.shape), aspect='equal')
  #plt.title ('init frame %d %s' % (start_f, img_filename))
  plt.set_cmap ('gray')
  #plt.colorbar()

  def drawbox(image, i):
    i = int(i)
    numOfPersoninFrame = len(frames[i])



    if numOfPersoninFrame > 0:
        for p in range(numOfPersoninFrame):
            hight = frames[i][p][1]
            width = frames[i][p][2]
            x = frames[i][p][3]
            y = frames[i][p][4]

            for k in range(int(hight)):
                image[int(y)+k, int(x)] = 255
            for k in range(int(hight)):
                image[int(y)+k, int(x)+int(width)] = 255
            for k in range(int(width)):
                image[int(y), int(x)+k] = 255
            for k in range(int(width)):
                image[int(y)+int(hight), int(x)+k] = 255

  # i is the i-th figure update starting from 0
  # Set frame # f = start_f + i so video plays from the specified start frame.
  def updatefig(i):
    f = start_f + i
    fdigs = digits_format % (f)
    img_filename = '%s%s.%s' % (imgseq_prefix, fdigs, imgseq_ftype)
    img_fullpath = imgseq_path+'/'+img_filename
    print '  updatefig reading %s' % (img_fullpath)
    try:
      image = io.imread (img_fullpath)
      drawbox(image, i)
    except IOError:
      print 'updatefig Error in reading %s' % (img_fullpath)
      return False

    plt_im.set_array(image)
    plt.title ('%d frame %d %s' % (i, f, img_filename))
    return plt_im

  #total video frame is 100
  #FPS: interval in milliseconds.
  #ani = animation.FuncAnimation(fig, updatefig, frames=100, interval=10, blit=True)
  print 'Video playback fps = %.1f ms' % (playback_fps)
  ani = animation.FuncAnimation(fig, updatefig, frames=n_f, interval=playback_fps, blit=False)
  plt.show()


if 1:
    # Play back original image sequence
    imgseg_showbox(IMGSEQ_PATH, IMGSEQ_PREFIX, IMGSEQ_NDIGITS, IMGSEQ_FTYPE, START_F, END_F, PLAYBACK_FPS, frames)